#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SIMION 无GUI自动化脚本
遍历频率列表，执行SIMION fly命令，解析输出并生成Excel文件
"""

import os
import sys
import subprocess
import re
import time
import signal
import threading
import math
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed

# 设置环境变量确保正确的编码
os.environ['PYTHONIOENCODING'] = 'utf-8'

# 全局停止标志
stop_requested = False

def signal_handler(signum, frame):
    """信号处理器，用于优雅停止"""
    global stop_requested
    stop_requested = True
    script_output("收到停止信号，正在安全退出...")

# 配置Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
if parent_dir not in sys.path:
    sys.path.insert(0, parent_dir)

# 导入配置（自动配置Python路径）
try:
    from settings import BASE_DIR
    from store import get_param
    from utils.logger import logger
    from utils.csv_utils import CSVUtils
    from utils.simion_workspace import SimionWorkspaceManager
    print("模块导入成功")
except Exception as e:
    print(f"模块导入失败: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)

# 创建一个简单的输出函数，用于脚本输出到主窗口（不包含时间戳）
def script_output(message):
    """脚本输出函数，直接输出到stdout，不包含时间戳"""
    print(message)
    sys.stdout.flush()


# 全局文件锁
lua_file_lock = threading.Lock()

def update_lua_frequency(lua_path, frequency, thread_id=None):
    """
    更新Lua文件中的频率参数（修复多线程文件竞争问题）
    
    Args:
        lua_path (str): Lua文件路径
        frequency (float): 新的频率值
        thread_id (str): 线程ID（用于日志显示）
        
    Returns:
        bool: 更新是否成功
    """
    try:
        if not os.path.exists(lua_path):
            script_output(f"Lua文件不存在: {lua_path}")
            return False
        

        with open(lua_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 使用精确的正则表达式替换频率值
        patterns = [
            r'(adjustable\s+_frequency\s*=\s*)\d+\.?\d*',  # 匹配整数和小数
            r'(adjustable\s+_frequency\s*=\s*)\d+',        # 只匹配整数
            r'(_frequency\s*=\s*)\d+\.?\d*',              # 更宽松的匹配
        ]
        
        new_content = content
        pattern_used = None
        frequency_updated = False
        
        for pattern in patterns:
            updated = re.sub(pattern, f'\\g<1>{int(frequency)}', new_content)
            if updated != new_content:
                new_content = updated
                pattern_used = pattern
                frequency_updated = True
                break
        
        # 检查是否找到频率参数（即使没有变化也算成功）
        if not frequency_updated:
            # 检查是否已经包含目标频率
            if f'_frequency = {int(frequency)}' in content:
                frequency_updated = True
            else:
                return False
        
        # 如果频率已经是目标值，直接返回成功
        if not frequency_updated:
            return True
        
        # 写回文件（仅锁写操作）
        with lua_file_lock:
            with open(lua_path, 'w', encoding='utf-8') as f:
                f.write(new_content)
        
        return True
        
    except Exception as e:
        return False

def process_single_frequency(args):
    """处理单个频率的线程函数"""
    simion_path, iob_path, lua_path, csv_save_path, frequency, index, total, workspace_manager = args
    
    # 获取当前线程ID
    thread_id = threading.current_thread().name
    
    try:
        
        if workspace_manager:
            # 多线程模式：使用工作空间，实现线程池复用
            # 计算实际使用的线程目录编号（循环复用）
            thread_count = workspace_manager.thread_count
            run_id = ((index - 1) % thread_count) + 1
            thread_workspace = workspace_manager.get_thread_workspace(run_id)
            thread_iob_path = thread_workspace['iob_path']
            thread_work_dir = thread_workspace['work_dir']
            
            # 获取线程专用的Lua文件路径（工作空间管理器已在创建时复制）
            thread_lua_path = os.path.join(thread_work_dir, os.path.basename(lua_path))
            
        else:
            # 单线程模式：直接使用原始文件
            thread_iob_path = iob_path
            thread_work_dir = os.path.dirname(iob_path)
            thread_lua_path = lua_path
        
        # 执行SIMION fly命令
        success, output_text, splats = run_simion_fly(simion_path, thread_iob_path, thread_lua_path, frequency, thread_id, thread_work_dir)
        
        if success and splats:
            # 计算发射度
            emittance_data = calculate_emittance(splats)
            
            # 计算扩展参数
            extended_params = calculate_extended_parameters(splats)
            
            # 使用新的CSV创建函数，包含发射度信息和扩展参数
            try:
                csv_path = CSVUtils.create_simion_csv_with_extended_data(
                    splats, frequency, csv_save_path, emittance_data, extended_params
                )
                
                if csv_path:
                    # 输出发射度信息
                    script_output(f"发射度计算结果:")
                    script_output(f"  RMS发射度 (x): {emittance_data['rms_x']:.2e}")
                    script_output(f"  RMS发射度 (y): {emittance_data['rms_y']:.2e}")
                    script_output(f"  几何发射度 (x): {emittance_data['geometric_x']:.2e}")
                    script_output(f"  几何发射度 (y): {emittance_data['geometric_y']:.2e}")
                    script_output(f"  规范化发射度 (x): {emittance_data['normalized_x']:.2e}")
                    script_output(f"  规范化发射度 (y): {emittance_data['normalized_y']:.2e}")
                    script_output(f"凸包面积: {extended_params['convex_hull_area']:.2e}")
                    script_output(f"CSV文件已保存: {os.path.basename(csv_path)} ({len(splats)} 个数据)")
                    return True, frequency, f"成功处理频率 {frequency} Hz"
                else:
                    return False, frequency, f"保存CSV文件失败"
                    
            except Exception as e:
                return False, frequency, f"处理CSV文件失败: {e}"
        else:
            return False, frequency, f"SIMION命令执行失败"
            
    except Exception as e:
        return False, frequency, f"处理错误: {e}"

def calculate_emittance(splats):
    """
    计算发射度（RMS发射度、几何发射度、规范化发射度）
    
    Args:
        splats: 包含(x, y, z, mass, vx, vy, vz, ke, t, charge)元组的列表
        
    Returns:
        dict: 包含各种发射度值的字典
    """
    if len(splats) < 2:
        return {
            'rms_x': 0.0, 'rms_y': 0.0,
            'geometric_x': 0.0, 'geometric_y': 0.0,
            'normalized_x': 0.0, 'normalized_y': 0.0
        }
    
    # 提取x, y, vx, vy数据
    x_data = [splat[0] for splat in splats]  # x坐标
    y_data = [splat[1] for splat in splats]  # y坐标
    vx_data = [splat[4] for splat in splats]  # vx速度
    vy_data = [splat[5] for splat in splats]  # vy速度
    
    # 计算x' = vx/vz (假设z方向为主轴，vz为轴向速度)
    # 这里我们使用vz的平均值作为参考
    vz_data = [splat[6] for splat in splats]  # vz速度
    vz_avg = np.mean(vz_data) if vz_data else 1.0
    
    x_prime = [vx / vz_avg for vx in vx_data]  # x' = vx/vz
    y_prime = [vy / vz_avg for vy in vy_data]  # y' = vy/vz
    
    # 计算RMS发射度
    def calculate_rms_emittance(pos_data, angle_data):
        if len(pos_data) < 2:
            return 0.0
        
        # 计算统计量
        pos_mean = np.mean(pos_data)
        angle_mean = np.mean(angle_data)
        
        # 计算方差和协方差
        pos_var = np.var(pos_data)
        angle_var = np.var(angle_data)
        pos_angle_cov = np.cov(pos_data, angle_data)[0, 1]
        
        # RMS发射度公式: sqrt(<x²><x'²> - <xx'>²)
        rms_emittance = math.sqrt(max(0, pos_var * angle_var - pos_angle_cov**2))
        return rms_emittance
    
    # 计算x和y方向的RMS发射度
    rms_x = calculate_rms_emittance(x_data, x_prime)
    rms_y = calculate_rms_emittance(y_data, y_prime)
    
    # 几何发射度 = 4 * RMS发射度
    geometric_x = 4 * rms_x
    geometric_y = 4 * rms_y
    
    # 规范化发射度计算（需要β和γ因子）
    # 这里使用简化的方法，假设粒子能量相对较低
    ke_data = [splat[7] for splat in splats]  # 动能数据
    ke_avg = np.mean(ke_data) if ke_data else 1000.0  # 平均动能，单位eV
    
    # 简化的β计算（假设非相对论情况）
    # β = v/c，这里使用平均速度估算
    v_avg = math.sqrt(np.mean([vx**2 + vy**2 + vz**2 for vx, vy, vz in zip(vx_data, vy_data, vz_data)]))
    c = 299792458  # 光速，m/s
    beta = min(v_avg * 1e6 / c, 0.99)  # 转换为m/s并限制β值
    gamma = 1.0 / math.sqrt(1 - beta**2) if beta < 0.99 else 10.0
    
    # 规范化发射度 = β * γ * 几何发射度
    normalized_x = beta * gamma * geometric_x
    normalized_y = beta * gamma * geometric_y
    
    return {
        'rms_x': rms_x,
        'rms_y': rms_y,
        'geometric_x': geometric_x,
        'geometric_y': geometric_y,
        'normalized_x': normalized_x,
        'normalized_y': normalized_y
    }

def calculate_velocity_magnitude(vx, vy, vz):
    """计算速度模长"""
    return math.sqrt(vx**2 + vy**2 + vz**2)

def calculate_momentum_components(mass, vx, vy, vz):
    """计算动量分量 p_i = m * v_i"""
    return mass * vx, mass * vy, mass * vz

def calculate_angles(vx, vy, vz):
    """计算角度 θx = arctan(vx/vz), θy = arctan(vy/vz)"""
    if abs(vz) < 1e-10:  # 避免除零错误
        vz = 1e-10 if vz >= 0 else -1e-10
    theta_x = math.atan(vx / vz)
    theta_y = math.atan(vy / vz)
    return theta_x, theta_y

def calculate_convex_hull_area(x_data, y_data):
    """
    计算凸包面积
    使用Graham扫描算法或scipy的ConvexHull
    """
    if len(x_data) < 3:
        return 0.0
    
    try:
        from scipy.spatial import ConvexHull
        # 创建点集
        points = np.column_stack((x_data, y_data))
        # 计算凸包
        hull = ConvexHull(points)
        return hull.volume  # 对于2D，volume就是面积
    except ImportError:
        # 如果没有scipy，使用简单的凸包算法
        return simple_convex_hull_area(x_data, y_data)

def simple_convex_hull_area(x_data, y_data):
    """
    简单的凸包面积计算（不使用scipy）
    使用Graham扫描算法的简化版本
    """
    if len(x_data) < 3:
        return 0.0
    
    # 找到最底部的点（y最小，如果相同则x最小）
    points = list(zip(x_data, y_data))
    bottom_point = min(points, key=lambda p: (p[1], p[0]))
    
    # 按极角排序
    def polar_angle(p):
        dx = p[0] - bottom_point[0]
        dy = p[1] - bottom_point[1]
        return math.atan2(dy, dx)
    
    # 移除重复点并排序
    unique_points = list(set(points))
    if len(unique_points) < 3:
        return 0.0
    
    sorted_points = sorted(unique_points, key=polar_angle)
    
    # Graham扫描算法
    hull = []
    for point in sorted_points:
        while len(hull) > 1 and cross_product(hull[-2], hull[-1], point) <= 0:
            hull.pop()
        hull.append(point)
    
    # 计算凸包面积（使用鞋带公式）
    if len(hull) < 3:
        return 0.0
    
    area = 0.0
    n = len(hull)
    for i in range(n):
        j = (i + 1) % n
        area += hull[i][0] * hull[j][1]
        area -= hull[j][0] * hull[i][1]
    return abs(area) / 2.0

def cross_product(o, a, b):
    """计算向量叉积 (a-o) × (b-o)"""
    return (a[0] - o[0]) * (b[1] - o[1]) - (a[1] - o[1]) * (b[0] - o[0])

def calculate_extended_parameters(splats):
    """
    计算扩展参数：速度模长、动量分量、角度、凸包面积
    
    Args:
        splats: 包含(x, y, z, mass, vx, vy, vz, ke, t, charge)元组的列表
        
    Returns:
        dict: 包含所有计算结果的字典
    """
    if len(splats) < 1:
        return {
            'velocity_magnitude': [],
            'momentum_x': [], 'momentum_y': [], 'momentum_z': [],
            'theta_x': [], 'theta_y': [],
            'convex_hull_area': 0.0
        }
    
    # 提取数据
    x_data = [splat[0] for splat in splats]
    y_data = [splat[1] for splat in splats]
    z_data = [splat[2] for splat in splats]
    mass_data = [splat[3] for splat in splats]
    vx_data = [splat[4] for splat in splats]
    vy_data = [splat[5] for splat in splats]
    vz_data = [splat[6] for splat in splats]
    
    # 计算速度模长
    velocity_magnitude = [calculate_velocity_magnitude(vx, vy, vz) 
                         for vx, vy, vz in zip(vx_data, vy_data, vz_data)]
    
    # 计算动量分量
    momentum_x = [mass * vx for mass, vx in zip(mass_data, vx_data)]
    momentum_y = [mass * vy for mass, vy in zip(mass_data, vy_data)]
    momentum_z = [mass * vz for mass, vz in zip(mass_data, vz_data)]
    
    # 计算角度
    theta_x = [calculate_angles(vx, vy, vz)[0] for vx, vy, vz in zip(vx_data, vy_data, vz_data)]
    theta_y = [calculate_angles(vx, vy, vz)[1] for vx, vy, vz in zip(vx_data, vy_data, vz_data)]
    
    # 计算凸包面积
    convex_hull_area = calculate_convex_hull_area(x_data, y_data)
    
    return {
        'velocity_magnitude': velocity_magnitude,
        'momentum_x': momentum_x,
        'momentum_y': momentum_y,
        'momentum_z': momentum_z,
        'theta_x': theta_x,
        'theta_y': theta_y,
        'convex_hull_area': convex_hull_area
    }

def parse_simion_output(output_text):
    """
    解析SIMION输出，提取Splat数据
    
    Args:
        output_text (str): SIMION命令的输出文本
        
    Returns:
        list: 包含(x, y, z, mass, vx, vy, vz, ke, t, charge)元组的列表
    """
    splats = []
    
    # SIMION标准Splat输出格式（包含完整信息，锚定整行，避免部分匹配导致列错位）
    splat_pattern = r'^Splat at x=([\d.-]+) mm, y=([\d.-]+) mm, z=([\d.-]+) mm, mass=([\d.-]+) amu, vx=([\d.-]+) mm/us, vy=([\d.-]+) mm/us, vz=([\d.-]+) mm/us, KE=([\d.-]+) eV, t=([\d.-]+) us, charge=([\d.-]+)$'
    
    # 解析Splat数据
    lines = [ln.strip() for ln in output_text.split('\n') if ln.strip()]
    for line in lines:
        match = re.search(splat_pattern, line)
        if match:
            try:
                x = float(match.group(1))
                y = float(match.group(2))
                z = float(match.group(3))
                mass = float(match.group(4))
                vx = float(match.group(5))
                vy = float(match.group(6))
                vz = float(match.group(7))
                ke = float(match.group(8))
                t = float(match.group(9))
                charge = float(match.group(10))
                
                splats.append((x, y, z, mass, vx, vy, vz, ke, t, charge))
                    
            except ValueError as e:
                continue
    
    
    # 如果没有解析到数据，检查输出格式
    if len(splats) == 0:
        splat_lines = [line for line in lines if 'splat' in line.lower()]
        if splat_lines:
            script_output("警告：检测到Splat输出，但格式不是期望的标准格式")
            script_output("请检查SIMION配置，确保输出标准格式的Splat数据")
    
    return splats


def run_simion_fly(simion_path, iob_path, lua_path, frequency, thread_id=None, work_dir=None):
    """
    执行SIMION fly命令
    
    Args:
        simion_path (str): SIMION可执行文件路径
        iob_path (str): IOB文件路径
        lua_path (str): Lua文件路径
        frequency (float): 当前频率值
        
    Returns:
        tuple: (success, output_text, splats)
    """
    try:
        # 先更新Lua文件中的频率参数
        if not update_lua_frequency(lua_path, frequency, thread_id):
            return False, "Lua文件更新失败", []
        
        # 构建命令
        cmd = [simion_path, '--nogui', '--quiet', 'fly', iob_path]
        
        # 执行命令（无超时限制）
        # 使用独立的工作目录避免文件冲突
        cwd = work_dir if work_dir else os.path.dirname(iob_path)
        
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            encoding='gbk',  # 使用GBK编码处理中文
            errors='replace',  # 处理编码错误
            cwd=cwd  # 设置工作目录
        )
        
        # 安全地合并输出，处理编码问题
        try:
            output_text = result.stdout + result.stderr
        except UnicodeDecodeError:
            # 如果仍有编码问题，使用错误替换模式
            output_text = (result.stdout.encode('utf-8', errors='replace').decode('utf-8') + 
                          result.stderr.encode('utf-8', errors='replace').decode('utf-8'))
        
        if result.returncode == 0:
            splats = parse_simion_output(output_text)
            return True, output_text, splats
        else:
            return False, output_text, []
            
    except Exception as e:
        return False, str(e), []


def main():
    """主函数"""
    start_time = time.time()  # 记录开始时间
    
    try:
        script_output("=== 开始SIMION无GUI自动化 ===")
        
        # 获取参数
        script_output("正在获取参数...")
        simion_path = get_param('simion_path')
        iob_path = get_param('iob_path')
        lua_path = get_param('lua_path')
        csv_save_path = get_param('csv_save_path')
        frequency_list = get_param('_frequency_list', [])
        thread_count = get_param('thread_count', 2)
        script_output("参数获取完成")
        
        # 验证参数
        if not simion_path:
            script_output("SIMION路径未设置")
            return
        
        if not iob_path:
            script_output("IOB路径未设置")
            return
        
        if not lua_path:
            script_output("Lua路径未设置")
            return
            
        if not csv_save_path:
            script_output("CSV保存路径未设置")
            return
            
        if not frequency_list:
            script_output("频率列表为空")
            return
        
        script_output(f"频率列表长度: {len(frequency_list)}")
        script_output(f"线程数量: {thread_count}")
        
        # 确保CSV保存目录存在
        os.makedirs(csv_save_path, exist_ok=True)
        
        # 设置SIMION工作空间（避免多线程文件冲突）
        workspace_manager = SimionWorkspaceManager(iob_path, thread_count)
        if not workspace_manager.setup_workspace():
            script_output("工作空间设置失败，使用单线程模式")
            thread_count = 1
            workspace_manager = None
        else:
            # 获取工作空间信息
            workspace_info = workspace_manager.get_workspace_info()
            electrode_count = workspace_info.get('electrode_count', 0)
            pa_extensions = workspace_info.get('pa_temp_extensions', [])
            
        
        # 验证工作空间文件是否准备就绪
        if workspace_manager:
            try:
                script_output("正在验证工作空间文件...")
                workspace_manager.verify_all_thread_files_ready()
                script_output("工作空间验证成功")
            except FileNotFoundError as e:
                script_output(f"工作空间验证失败: {e}")
                return
        else:
            script_output("使用单线程模式，跳过工作空间验证")
        
        # 准备线程参数
        total_frequencies = len(frequency_list)
        if workspace_manager:
            # 多线程模式
            thread_args = [
                (simion_path, iob_path, lua_path, csv_save_path, freq, i, total_frequencies, workspace_manager)
                for i, freq in enumerate(frequency_list, 1)
            ]
        else:
            # 单线程模式
            thread_args = [
                (simion_path, iob_path, lua_path, csv_save_path, freq, i, total_frequencies, None)
                for i, freq in enumerate(frequency_list, 1)
            ]
        
        # 设置信号处理器
        signal.signal(signal.SIGTERM, signal_handler)
        signal.signal(signal.SIGINT, signal_handler)
        
        # 显示线程池配置信息
        script_output(f"准备启动线程池: {thread_count} 个线程, {total_frequencies} 个任务")
        
        # 使用线程池处理
        success_count = 0
        failed_count = 0
        
        with ThreadPoolExecutor(max_workers=thread_count) as executor:
            # 提交所有任务
            future_to_freq = {
                executor.submit(process_single_frequency, args): args[4] 
                for args in thread_args
            }
            
            # 处理完成的任务
            for future in as_completed(future_to_freq):
                if stop_requested:
                    # 取消未完成的任务
                    for f in future_to_freq:
                        f.cancel()
                    break
                
                frequency = future_to_freq[future]
                try:
                    success, freq, message = future.result()
                    if success:
                        success_count += 1
                    else:
                        failed_count += 1
                except Exception as e:
                    failed_count += 1
        
        # 计算运行时间
        end_time = time.time()
        total_time = end_time - start_time
        
        # 格式化时间显示
        hours = int(total_time // 3600)
        minutes = int((total_time % 3600) // 60)
        seconds = int(total_time % 60)
        
        if hours > 0:
            time_str = f"{hours}小时{minutes}分钟{seconds}秒"
        elif minutes > 0:
            time_str = f"{minutes}分钟{seconds}秒"
        else:
            time_str = f"{seconds}秒"
        
        # 输出统计信息
        script_output("=== SIMION无GUI自动化完成 ===")
        script_output(f"成功处理: {success_count} 个频率")
        script_output(f"失败处理: {failed_count} 个频率")
        script_output(f"总计: {success_count + failed_count} 个频率")
        script_output(f"总运行时间: {time_str} ({total_time:.2f}秒)")
        
        # 清理工作空间（根据keep_workspace参数决定）
        if workspace_manager:
            keep_workspace = get_param('keep_workspace', False)
            if keep_workspace:
                workspace_manager.cleanup_workspace(keep_workspace=True)
            else:
                workspace_manager.cleanup_workspace(keep_workspace=False)
        
    except Exception as e:
        # 即使出错也显示运行时间
        end_time = time.time()
        total_time = end_time - start_time
        hours = int(total_time // 3600)
        minutes = int((total_time % 3600) // 60)
        seconds = int(total_time % 60)
        
        if hours > 0:
            time_str = f"{hours}小时{minutes}分钟{seconds}秒"
        elif minutes > 0:
            time_str = f"{minutes}分钟{seconds}秒"
        else:
            time_str = f"{seconds}秒"
            
        script_output(f"自动化过程中发生错误: {e}")
        script_output(f"错误类型: {type(e).__name__}")
        import traceback
        script_output(f"错误详情: {traceback.format_exc()}")
        script_output(f"运行时间: {time_str} ({total_time:.2f}秒)")


if __name__ == "__main__":
    main()
